1792D - Fixed Prefix Permutations - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;

struct TrieNode {
    struct TrieNode* ch[11];
    TrieNode() {
        for (int i = 0; i < 11; i++) {
            ch[i] = nullptr;
        }
    }

    ~TrieNode() {
        for (int i = 0; i < 11; i++) {
            if (ch[i] != nullptr) {
                delete ch[i];
            }
        }
    }
};

TrieNode* root;

void add(TrieNode* node, const vector<int>& a) {
    for(int i = 0; i < a.size(); i++) {
        int id = a[i];
        if (node->ch[id] == nullptr) {
            node->ch[id] = new TrieNode;
        }

        node = node->ch[id];
    }
}

int check(TrieNode* node, const vector<int>& a) {
    int cnt = 0;
    for (int i = 0; i < a.size(); i++) {
        int id = a[i];
        if (node->ch[id] != nullptr) {
            node = node->ch[id];
            ++cnt;
        } else {
            return cnt;
        }
    }

    return cnt;
}

vector<int> div(vector<int>& a) {
    vector<int> res(a.size(), 0);
    for (int i = 0; i < a.size(); i++) {
        int id = a[i] - 1;
        res[id] = i+1;
    }

    return res;
}

void solve() {
    int n, m;
    scanf("%d%d", &n, &m);
    root = new TrieNode;
    vector<vector<int>> perms(n, vector<int>(m, 0));
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < m; j++) {
            scanf("%d", &perms[i][j]);
        }

        vector<int> c = div(perms[i]);
        add(root, c);
    }

    int ans = 0;
    for (int i = 0; i < n; i++) {
        int t = check(root, perms[i]);
        printf(i == 0 ? "%d":" %d", t);
    }

    printf("\n");
    delete root;
}

int main() {
    int t;
    scanf("%d", &t);
    while (t--) {
        solve();
    }
    return 0;
}